from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal


def main():

    seq_len = 15
    redundancy = 5
    diverse_set_capacity = 5

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--seq_len", type=int, default=8)
    parser.add_argument("--redundancy", type=int, default=3)
    parser.add_argument("--diverse_set_capacity", type=int, default=5)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    redundancy = args.redundancy
    seq_len = args.seq_len
    diverse_set_capacity = args.diverse_set_capacity

    print("redun: ", redundancy)
    print("seq_len", seq_len)

    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    rnn_limit_rings_file_name = "./logs/rnn_limit_rings_of_best_estimation_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
    # 载入 npz 文件
    rnn_limit_rings_of_best_estimation_file = np.load(rnn_limit_rings_file_name)

    # 获取 npz 文件中的所有对象名称
    matrix_names = rnn_limit_rings_of_best_estimation_file.files

    rnn_limit_rings_of_best_estimation = []

    # 遍历对象名称，访问和操作每个矩阵对象
    for name in matrix_names:
        matrix = rnn_limit_rings_of_best_estimation_file[name]
        # 在这里进行对矩阵对象的操作
        # 例如，打印矩阵的形状
        # print(f"Matrix '{name}' shape: {matrix.shape}")
        rnn_limit_rings_of_best_estimation.append(matrix)

    # 求 rnn_limit_rings_of_best_estimation 的中心位置序列
    rnn_limit_rings_of_best_estimation_center = []
    for i in range(len(rnn_limit_rings_of_best_estimation)):
        rnn_limit_rings_of_best_estimation_center.append(np.mean(rnn_limit_rings_of_best_estimation[i], axis=(0,1)))
    rnn_limit_rings_of_best_estimation_center = np.array(rnn_limit_rings_of_best_estimation_center)
    print("rnn_limit_rings_of_best_estimation_center.shape: ", rnn_limit_rings_of_best_estimation_center.shape)
    
    
    file_name = "obs_data_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
    obs_file = np.load("./logs/" + file_name)
    obs_data = obs_file["obs_data"]
    diverse_set_trajectoies = obs_file["diverse_set_trajectoies"]
    diverse_set_actions = obs_file["diverse_set_actions"]

    # 将所有轨迹的起点都对齐到原点
    for i in range(diverse_set_trajectoies.shape[0]):
        diverse_set_trajectoies[i] = diverse_set_trajectoies[i] - diverse_set_trajectoies[i][0]

    # 将 diverse_set_trajectoies 的后两维合并成一维, 例如 (100, 8, 2) 变成 (100, 16)
    diverse_set_trajectoies_2d = diverse_set_trajectoies.reshape(-1, seq_len * 2)

    # 对 diverse_set_trajectoies_2d 求取自相似矩阵
    diverse_set_trajectoies_2d_self_similarity = np.zeros((diverse_set_trajectoies_2d.shape[0], diverse_set_trajectoies_2d.shape[0]))
    for i in range(diverse_set_trajectoies_2d.shape[0]):
        for j in range(diverse_set_trajectoies_2d.shape[0]):
            diverse_set_trajectoies_2d_self_similarity[i][j] = np.linalg.norm(diverse_set_trajectoies_2d[i] - diverse_set_trajectoies_2d[j])

    # 对 rnn_limit_rings_of_best_estimation_center 求取自相似矩阵
    rnn_limit_rings_of_best_estimation_center_self_similarity = np.zeros((rnn_limit_rings_of_best_estimation_center.shape[0], rnn_limit_rings_of_best_estimation_center.shape[0]))
    for i in range(rnn_limit_rings_of_best_estimation_center.shape[0]):
        for j in range(rnn_limit_rings_of_best_estimation_center.shape[0]):
            rnn_limit_rings_of_best_estimation_center_self_similarity[i][j] = np.linalg.norm(rnn_limit_rings_of_best_estimation_center[i] - rnn_limit_rings_of_best_estimation_center[j])
    
    print("shape of diverse_set_trajectoies_2d: ", diverse_set_trajectoies_2d.shape)
    print("shape of rnn_limit_rings_of_best_estimation_center: ", rnn_limit_rings_of_best_estimation_center.shape)

    # 将所有 diverse_set_trajectoies 绘制成plot连线图
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111)
    for i in range(diverse_set_trajectoies.shape[0]):
        ax.plot(diverse_set_trajectoies[i,:,0], diverse_set_trajectoies[i,:,1], '-')
    plt.legend()
    plt.show()

    # 将两个自相似矩阵合并在同一个视图中,使用两个subplot进行可视化
    fig = plt.figure(figsize=(10, 10))
    ax1 = fig.add_subplot(121)
    ax1.imshow(diverse_set_trajectoies_2d_self_similarity)
    ax2 = fig.add_subplot(122)
    ax2.imshow(rnn_limit_rings_of_best_estimation_center_self_similarity)
    # 添加子图标题以及颜色条
    ax1.set_title("diverse_set_trajectoies_2d_self_similarity")
    ax2.set_title("rnn_limit_rings_of_best_estimation_center_self_similarity")
    fig.colorbar(ax1.imshow(diverse_set_trajectoies_2d_self_similarity), ax=ax1)
    fig.colorbar(ax2.imshow(rnn_limit_rings_of_best_estimation_center_self_similarity), ax=ax2)
    plt.show()

    # 将两个自相似矩阵转换成两个向量，然后计算两个向量的pearson相关系数
    diverse_set_trajectoies_2d_self_similarity_vector = diverse_set_trajectoies_2d_self_similarity.reshape(-1)
    rnn_limit_rings_of_best_estimation_center_self_similarity_vector = rnn_limit_rings_of_best_estimation_center_self_similarity.reshape(-1)
    print("diverse_set_trajectoies_2d_self_similarity_vector.shape: ", diverse_set_trajectoies_2d_self_similarity_vector.shape)
    print("rnn_limit_rings_of_best_estimation_center_self_similarity_vector.shape: ", rnn_limit_rings_of_best_estimation_center_self_similarity_vector.shape)
    # 计算 pearson 相关系数
    pearson_correlation = np.corrcoef(diverse_set_trajectoies_2d_self_similarity_vector, rnn_limit_rings_of_best_estimation_center_self_similarity_vector)
    print("pearson_correlation: ", pearson_correlation)

    def matrix_similarity(A, B):
        mse = np.mean((A - B) ** 2)
        return mse
    
    # 计算两个自相似矩阵的 mse
    diverse_set_trajectoies_2d_self_similarity_norm = diverse_set_trajectoies_2d_self_similarity / np.max(diverse_set_trajectoies_2d_self_similarity)
    rnn_limit_rings_of_best_estimation_center_self_similarity_norm = rnn_limit_rings_of_best_estimation_center_self_similarity / np.max(rnn_limit_rings_of_best_estimation_center_self_similarity)
    mse = matrix_similarity(diverse_set_trajectoies_2d_self_similarity_norm, rnn_limit_rings_of_best_estimation_center_self_similarity_norm)
    print("mse: ", mse)

    def shen_similarity(A, B):
        # A和B是两个NxN邻接矩阵
        min_ab = np.minimum(A, B) 
        max_ab = np.maximum(A, B)
        denom = max_ab.sum()
        numer = min_ab.sum()
        return numer/denom
    
    # 计算两个自相似矩阵的 shen 相似度
    shen_similarity = shen_similarity(diverse_set_trajectoies_2d_self_similarity_norm, rnn_limit_rings_of_best_estimation_center_self_similarity_norm)
    print("shen_similarity: ", shen_similarity)

    # # 对 diverse_set_trajectoies_2d 进行 t-SNE 降维并且可视化
    # tsne = TSNE(n_components=2, random_state=0)
    # diverse_set_trajectoies_2d_tsne = tsne.fit_transform(diverse_set_trajectoies_2d)
    # print("diverse_set_trajectoies_2d_tsne.shape: ", diverse_set_trajectoies_2d_tsne.shape)
    
    # plt.figure(figsize=(10, 10))
    # plt.scatter(diverse_set_trajectoies_2d_tsne[:,0], diverse_set_trajectoies_2d_tsne[:,1], s=5)
    # plt.show()

    # # 对 rnn_limit_rings_of_best_estimation_center 进行 t-SNE 降维并且可视化
    # tsne = TSNE(n_components=2, random_state=0)
    # rnn_limit_rings_of_best_estimation_center_tsne = tsne.fit_transform(rnn_limit_rings_of_best_estimation_center)
    # print("rnn_limit_rings_of_best_estimation_center_tsne.shape: ", rnn_limit_rings_of_best_estimation_center_tsne.shape)

    # plt.figure(figsize=(10, 10))
    # plt.scatter(rnn_limit_rings_of_best_estimation_center_tsne[:,0], rnn_limit_rings_of_best_estimation_center_tsne[:,1], s=5)
    # plt.show()

    # 对 rnn_limit_rings_of_best_estimation_center 进行 PCA 并且进行3D可视化
    pca = PCA()
    rnn_limit_rings_of_best_estimation_center_pca = pca.fit_transform(rnn_limit_rings_of_best_estimation_center)
    print("rnn_limit_rings_of_best_estimation_center_pca.shape: ", rnn_limit_rings_of_best_estimation_center_pca.shape)

    # 打印出 pca 的所有主成分的方差比例
    print("pca.explained_variance_ratio_: ", pca.explained_variance_ratio_)

    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(rnn_limit_rings_of_best_estimation_center_pca[:,0], rnn_limit_rings_of_best_estimation_center_pca[:,1], rnn_limit_rings_of_best_estimation_center_pca[:,2], s=5)
    
    # 获取X轴的范围
    x_min, x_max = rnn_limit_rings_of_best_estimation_center_pca[:,0].min(), rnn_limit_rings_of_best_estimation_center_pca[:,0].max()
    x_min, x_max = x_min*1.5, x_max*1.5
    # 设置Y轴和Z轴的范围与X轴一致
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(x_min, x_max)
    ax.set_zlim(x_min, x_max)

    plt.show()


if __name__ == "__main__":
    main()